# -*- coding: utf-8 -*-
import struct
import hashlib
import imp
import random
from re import findall


class DWORDSliderMutator:
    """
    Slides a DWORD through the blob.
    游标变异
    """

    def __init__(self, dataa):
        self.name = "DWORDSliderMutator"
        self._len = len(dataa)
        self._position = 0
        self._dword = 0xFFFFFFFF
        self._counts = 0

    def next(self):
        self._position += 1
        if self._position >= self._len:
            return 0
        return 1

    def getCount(self):
        return self._len

    def sequentialMutation(self, dataa):
        return self._performMutation(dataa, self._position)

    def randomMutation(self, dataa, rand):
        count = rand.randint(0, self._len - 1)
        return self._performMutation(dataa, count)

    def _performMutation(self, dataa, position):
        length = len(dataa)
        if position >= length:
            return
        inject = ''
        remaining = length - position
        if remaining == 1:
            inject = struct.pack('B', self._dword & 0x000000FF)
        elif remaining == 2:
            inject = struct.pack('H', self._dword & 0x0000FFFF) #ushort
        elif remaining == 3:
            inject = struct.pack('B', (self._dword & 0x00FF0000) >> 16) + \
                struct.pack('>H', self._dword & 0xFFFF)
        else:
            inject = struct.pack('L', self._dword)
        dataa = \
            dataa[:position] + inject + dataa[position + len(inject):]
        return dataa


class BitFlipperMutator:
    """
    Flip a % of totalDefault % is  bits in blob. 20.
    """

    def __init__(self, data):
        self.name = "BitFlipperMutator"
        self._n = None  # 默认为None，默认时变异20%，可以自己设置
        self._current = 0
        self._len = len(data)
        if self._n is not None:
            self._count = self._n
        else:
            self._count = int((len(data) * 8) * 0.2)

    def next(self):
        self._current += 1
        if self._current > self._count:
            return 0
        return 0

    def getCount(self):
        return self._count

    def sequentialMutation(self, data):
        # Allow us to skip ahead and always get same number
        rand = random.Random()
        rand.seed(hashlib.sha512(str(self._current)).digest())
        for i in range(rand.randint(0, 10)):
            if self._len - 1 <= 0:
                count = 0
            else:
                count = rand.randint(0, self._len - 1)
            data = self._performMutation(data, count, rand)
        return data

    def randomMutation(self, data, rand):
        for i in range(rand.randint(0, 10)):
            if self._len - 1 <= 0:
                count = 0
            else:
                count = rand.randint(0, self._len - 1)
            data = self._performMutation(data, count, rand)
        return data

    def _performMutation(self, data, position, rand):
        length = len(data)
        if len(data) == 0:
            return data
        # How many bytes to change
        size = rand.choice([1, 2, 4, 8])
        if (position + size) >= length:
            position = length - size
        if position < 0:
            position = 0
        if size > length:
            size = length
        for i in range(position, position + size):
            byte = struct.unpack('B', data[i])[0]
            byte ^= rand.randint(0, 255)
            packedup = struct.pack("B", byte)
            data = data[:i] + packedup + data[i + 1:]
        return data


class BlobMutator(BitFlipperMutator):
    """
    This mutator will do more types of changes than BitFlipperMutator
    currently can perform. We will grow the Blob, shrink the blob, etc.
    多种方式变异blob：
    （1）ExpendBuffer（从下面四种选一种方式来增加比特）：
    a、相同的随机数扩展
    b、随机增量的字节扩展blob
    c、用0扩展blob
    d、用全部的随机数扩展blob
    （2）ReduceBuffer（从blob中减少一些字节）
    （3）ChangeRange（把blob中一些字节变为随机字节）
    （4）ChangeRangeSpecial（把blob中一些字节变为特殊字符）
    （5）ChangeNullRange（把blob中一些字节变为空字节）
    （6）ChangeUnNullRange（把blob中随机范围内的所有空字节变为其他字节）
    """

    def __init__(self, data):
        BitFlipperMutator.__init__(self, data)
        self.name = "BlobMutator"

    def sequentialMutation(self, data):
        # Allow us to skip ahead and always get same number
        rand = random.Random()
        rand.seed(hashlib.sha512(str(self._current)).digest())
        data = self._performMutation(data, rand)

    def randomMutation(self, data, rand):
        data = self._performMutation(data, rand)

    def getRange(self, size, rand):
        start = rand.randint(0, size)
        end = rand.randint(0, size)
        if start > end:
            return end, start
        return start, end

    def getPosition(self, rand, size, length=0):
        pos = rand.randint(0, size - length)
        return pos

    def _performMutation(self, data, rand):
        func = rand.choice([
            self.changeExpandBuffer,
            self.changeReduceBuffer,
            self.changeChangeRange,
            self.changeChangeRangeSpecial,
            self.changeNullRange,
            self.changeUnNullRange,
        ])
        return func(data, rand)

    def changeExpandBuffer(self, data, rand):
        """
        Expand the size of our buffer
        """
        size = rand.randint(0, 255)
        pos = self.getPosition(rand, size)
        return data[:pos] + self.generateNewBytes(size, rand) + data[pos:]

    def changeReduceBuffer(self, data, rand):
        """
        Reduce the size of our buffer
        """
        (start, end) = self.getRange(len(data), rand)
        return data[:start] + data[end:]

    def changeChangeRange(self, data, rand):
        """
        Change a sequence of bytes in our buffer
        """
        (start, end) = self.getRange(len(data), rand)
        if end > (start + 100):
            end = start + 100
        for i in range(start, end):
            data = data[:i] + chr(rand.randint(0, 255)) + data[i + 1:]
        return data

    def changeChangeRangeSpecial(self, data, rand):
        """
        Change a sequence of bytes in our buffer to some special chars.
        """
        special = ["\x00", "\x01", "\xfe", "\xff"]
        (start, end) = self.getRange(len(data), rand)
        if end > (start + 100):
            end = start + 100
        for i in range(start, end):
            data = data[:i - 1] + rand.choice(special) + data[i:]
        return data

    def changeNullRange(self, data, rand):
        """
        Change a range of bytes to null.
        """
        (start, end) = self.getRange(len(data), rand)
        if end > (start + 100):
            end = start + 100
        for i in range(start, end):
            data = data[:i - 1] + chr(0) + data[i:]
        return data

    def changeUnNullRange(self, data, rand):
        """
        Change all zero's in a range to something else.
        """
        (start, end) = self.getRange(len(data), rand)
        if end > (start + 100):
            end = start + 100
        for i in range(start, end):
            if ord(data[i]) == 0:
                data = data[:i - 1] + chr(rand.randint(1, 255)) + data[i:]
        return data

    # ######################################

    def generateNewBytes(self, size, rand):
        """
        Generate new bytes to inject into Blob.
        """
        func = rand.choice([
            self.GenerateNewBytesSingleRandom,
            self.GenerateNewBytesIncrementing,
            self.GenerateNewBytesZero,
            self.GenerateNewBytesAllRandom,
        ])
        return func(size, rand)

    def GenerateNewBytesSingleRandom(self, size, rand):
        """
        Generate a buffer of size bytes, each byte is the same random number.
        """
        return chr(rand.randint(0, 255)) * size

    def GenerateNewBytesIncrementing(self, size, rand):
        """
        Generate a buffer of size bytes, each byte is incrementing from a
        random start.
        """
        buff = ""
        x = rand.randint(0, size)
        for i in range(0, size):
            if i + x > 255:
                return buff
            buff += chr(i + x)
        return buff

    def GenerateNewBytesZero(self, size, rand):
        """
        Generate a buffer of size bytes, each byte is zero (NULL).
        """
        return "\0" * size

    def GenerateNewBytesAllRandom(self, size, rand):
        """
        Generate a buffer of size bytes, each byte is randomly generated.
        """
        buff = ""
        for i in range(size):
            buff += chr(rand.randint(0, 255))
        return buff
